import uos
from micropython import const

OTA_NONE		= const(0)
OTA_DOWNLOAD 	= const(1)
OTA_TIMEOUT		= const(0xFF)
FW_NAME			= '/data/firmware.bin'
OTA_PACK_SIZE	= const(64)

class OTA:
	def __init__(self, send, recv, crypt, retry):
		self._status = OTA_NONE
		self._send = send
		self._recv = recv
		self._crypt = crypt
		self._ver = 0
		self._size = 0
		self._tag = bytearray(16)
		self._pos = 0
		self._file = None
		self._retry = retry

	def _packGet(self, count, pos):
		data = bytearray(6)		# 长度增加到4byte（原来只有2byte，只能支持64KB）
		
		data[0] = 0x02
		data[1] = count

		data[2] = (pos >> 24) & 0xFF
		data[3] = (pos >> 16) & 0xFF
		data[4] = (pos >> 8) & 0xFF
		data[5] = pos & 0xFF

		return data

	def _saveFW(self, data):
		self._file.write(data[2:])
		size = (self._pos * OTA_PACK_SIZE) + len(data) - 2
		print("size = ", size, "all = ", self._size)
		ret = size >= self._size
		if (ret):
			self._file.flush()
			self._file.close()
		else:
			self._pos += 1

		return ret

	def getFirmware(self, fwInfo, count):
		# 判断固件版本和当前缓存是否一致
		_same = False
		if ( (fwInfo[0] == self._ver) and (fwInfo[1] == self._size) and (fwInfo[2] == self._tag)):
			# same
			_same = True

		if (not _same):
			self._ver = fwInfo[0]
			self._size = fwInfo[1]
			self._tag[:] = fwInfo[2]

			# 清除旧文件，打开新文件
			if (self._file != None):
				self._file.close()

			self._file = open(FW_NAME, 'wb')
			self._pos = 0

		# 请求数据，等待响应
		data = self._packGet(count, self._pos)
		buf = self._crypt.encrypt(data)

		for i in range(self._retry):
			self._send(buf)

			# 等待确认（超时重发）
			data = self._recv()

			ret = data[0]
			if (ret):
				break
		
		last = False
		if (ret):
			# 成功接收数据，进行处理
			last = self._saveFW(data[1])

		return ret, last

	def stop(self):
		self._status = OTA_NONE
		self._ver = 0
		self._size = 0
		self._pos = 0
		self._file.close()
		self._file = None
		return

	